# -*- coding: utf-8 -*-
import os
import random
import time
import configparser
from pathlib import Path
from datetime import datetime, timedelta
from tqdm import tqdm
import logging
import string

class TestFileGenerator:
    def __init__(self, config_path='file_generator_config.ini'):
        self.config = self._load_config(config_path)
        self.logger = self._setup_logger()
        self.ext_dist = self._parse_distribution(self.config['Settings']['extensions'])
        self.keyword_dist = self._parse_distribution(self.config['Settings']['keywords'])
        self.time_range = self._parse_time_range()

    def _load_config(self, config_path):
        config = configparser.ConfigParser()
        with open(config_path, 'r', encoding='utf-8') as f:
            config.read_file(f)
        return config

    def _setup_logger(self):
        logging.basicConfig(
            filename='file_generator.log',
            format='%(asctime)s - %(levelname)s - %(message)s',
            level=logging.INFO
        )
        return logging.getLogger()

    def _parse_distribution(self, raw_str):
        """解析扩展名/关键词的分布配置"""
        items = [item.split(':') for item in raw_str.split(', ')]
        return {item[0].strip(): int(item[1]) for item in items}

    def _parse_time_range(self):
        """解析时间范围配置"""
        start = datetime.strptime(self.config['Settings']['start_date'], "%Y-%m-%d")
        end = datetime.strptime(self.config['Settings']['end_date'], "%Y-%m-%d")
        return (start, end + timedelta(days=1))  # 包含结束日期的全天

    def _generate_filename(self):
        """生成包含随机关键词的文件名"""
        keywords = random.choices(
            list(self.keyword_dist.keys()),
            weights=list(self.keyword_dist.values()),
            k=random.randint(0, 2)  # 每个文件包含0-2个关键词
        )
        random_str = ''.join(random.choices(string.ascii_lowercase, k=5))
        return f"{random_str}_{'_'.join(keywords)}"

    def _generate_file_content(self):
        """生成随机文件内容（1KB左右）"""
        safe_printable = string.printable.replace('\x0b', '').replace('\x0c', '')
        content = [
            f"Timestamp: {datetime.now().isoformat()}",
            "Sample log content:",
            ''.join(random.choices(safe_printable, k=900))
        ]
        return '\n'.join(content)

    def _set_file_mtime(self, file_path, modify_time):
        """修改文件访问和修改时间"""
        timestamp = modify_time.timestamp()
        os.utime(file_path, (timestamp, timestamp))

    def generate_files(self):
        target_dir = Path(os.path.normpath(self.config['Settings']['target_directory']))
        target_dir.mkdir(exist_ok=True, parents=True)

        total_files = self.config.getint('Settings', 'file_count')
        time_start, time_end = self.time_range
        time_delta = (time_end - time_start).total_seconds()  # 正确计算时间差
        start_time = time.time()
        self.logger.info(f"Starting to generate {total_files} test files...")

        try:
            with tqdm(total=total_files, desc="Generating files") as pbar:
                for _ in range(total_files):
                    random_seconds = random.uniform(0, time_delta)  # 使用数值参数
                    file_time = time_start + timedelta(seconds=random_seconds)
                    
                    # 生成文件名和扩展名
                    ext = random.choices(
                        list(self.ext_dist.keys()),
                        weights=list(self.ext_dist.values())
                    )[0]
                    filename = f"{self._generate_filename()}{ext}"
                    
                    # 创建文件
                    file_path = target_dir / filename
                    try:
                        with open(file_path, 'w') as f:
                            f.write(self._generate_file_content())
                        
                        # 修改文件时间
                        self._set_file_mtime(file_path, file_time)
                        pbar.update(1)
                    except Exception as e:
                        self.logger.error(f"Error creating {file_path}: {str(e)}")

            print(f"\nSuccessfully generated {total_files} test files in {target_dir}")
            print(f"总耗时: {time.time() - start_time:.2f} 秒")
            self.logger.info("File generation completed successfully")

        except KeyboardInterrupt:
            print("\nGeneration interrupted. Cleaning up...")
            print(f"已运行时间: {time.time() - start_time:.2f} 秒")
            self.logger.warning("File generation interrupted by user")
            # 可选：删除已生成的文件

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='测试文件生成工具')
    parser.add_argument('--dir', help='覆盖配置中的目标目录')
    args = parser.parse_args()

    generator = TestFileGenerator()
    if args.dir:
        generator.config['Settings']['target_directory'] = args.dir
    generator.generate_files() 